from pypng import png
import numpy

def loadImage(filename):
    # loads an image to an int numpy array
    f = open(filename, 'rb')
    reader = png.Reader(file=f)
    w, h, pixels, metadata = reader.asDirect()
    result = numpy.array([x for x in pixels], int)
    f.close()
    
    result.shape = (h, w, 4)
    result = numpy.transpose(result, (1, 0, 2))
    return result

def saveImage(filename, imageData):
    # saves an image from an int numpy array
    nx, ny, nc = imageData.shape
    writeData = numpy.transpose(imageData, (1, 0, 2))
    writeData = numpy.reshape(writeData, (ny, nx*nc))
    
    greyscale = (nc < 3)
    alpha = (nc == 2 or nc == 4)
    
    f = open(filename, 'wb')
    writer = png.Writer(width=nx, height=ny, greyscale=greyscale, alpha=alpha)
    writer.write(f, writeData)
    f.close()
